import os, sys
from pylab import *
from scipy.ndimage import gaussian_filter1d

class PhotozSNeDisp:
    def __init__(self,  n_mock=None, sample_size=None, \
                        mock_sn_dir=None,\
                        sn_dist_outroot=None ):
        self.zmin = 10.
        self.zmax = 0.0
        
        if n_mock is not None:
            self.n_mock = n_mock
        else:
            print("n_mock = None")
            sys.exit(0)
            
        if sample_size is not None:
            self.sample_size = sample_size
        else:
            print("sample_size is None")
            sys.exit(0)
            
        if mock_sn_dir is not None:
            self.mock_sn_dir = mock_sn_dir
        else:
            print("mock_sn_dir is None")
            sys.exit(0)
            
        if sn_dist_outroot is not None:
            self.sn_dist = sn_dist_outroot
        else:
            print("sn_dist_outroot is None")
            sys.exit(0)

        self.get_zmin_zmax()

    def get_zmin_zmax(self):
        print("==> searching zmin and zmax")
        cnt = 0
        while cnt < self.n_mock:
            sn = loadtxt(os.path.join(self.mock_sn_dir,'MOCK_'+str(cnt+1)+'.txt'))
            if self.zmin > sn[:,0].min():
                self.zmin = sn[:,0].min()

            if self.zmax < sn[:,0].max():
                self.zmax = sn[:,0].max()
            cnt += 1

        print '='*50
        print '--> searched zmin = ', self.zmin
        print '--> searched zmax = ', self.zmax
        print '='*50


    def estimate_disp(self):
        mean_mu = []
        disp_mu = []
        nbin = 50
        zbin = linspace(self.zmin,self.zmax,nbin+1)
        cnt = 0
        while cnt < self.n_mock:
            mock_file = os.path.join(self.mock_sn_dir,'MOCK_'+str(cnt+1)+'.txt')
            print '--> processing %s'%(mock_file)
            sn = loadtxt(mock_file)
            mean_mux = zeros(nbin)
            disp_mux = zeros(nbin)
            num_sn = zeros(nbin)

            for i in range(nbin):
                id1 = sn[:,0] >= zbin[i]
                id2 = sn[:,0] < zbin[i+1]
                idx = id1*id2
                mean_mux[i] = mean(sn[idx,1])
                disp_mux[i] = std(sn[idx,1])
                num_sn[i] = sum(idx)

            mean_mu.append(mean_mux)
            disp_mu.append(disp_mux)
            cnt += 1

        mean_mu = array(mean_mu)
        disp_mu = array(disp_mu)

        zbin = (zbin[1:]+zbin[:-1])/2.0
        zbin = zbin.reshape(len(zbin),1)
        mean_disp_mu = disp_mu.mean(axis=0)
        disp_mean_mu = mean_mu.std(axis=0)

        mean_disp_mu_smoothed = gaussian_filter1d(mean_disp_mu,sigma=2.0)
        disp_mean_mu_smoothed = gaussian_filter1d(disp_mean_mu,sigma=2.0)

        # do gaussian smoothing
        mean_disp_mu_smoothed = mean_disp_mu_smoothed.reshape(len(mean_disp_mu_smoothed),1)
        disp_mean_mu_smoothed = disp_mean_mu_smoothed.reshape(len(disp_mean_mu_smoothed),1)

        # save estimated variances into disk
        root1 = 'intrinsic_total_dispersion_'+self.sn_dist+'.txt'
        root2 = 'sample_variance_of_mean_mu_'+self.sn_dist+'.txt'
        root3 = 'TOTAL_VAR_'+self.sn_dist+'.txt'
        savetxt(root1, hstack((zbin,mean_disp_mu_smoothed)),fmt='%6.3f', delimiter=' ')
        savetxt(root2, hstack((zbin,disp_mean_mu_smoothed)),fmt='%6.3f', delimiter=' ')
        savetxt(root3, hstack((zbin,disp_mean_mu_smoothed+mean_disp_mu_smoothed)),fmt='%6.3f', delimiter=' ')

#        d1 = loadtxt(root1)
#        d2 = loadtxt(root2)
#        plot(d1[:,0],d1[:,1],'r-', label='intrinsic total dispersion')
#        plot(d2[:,0],d2[:,1],'b-', label='sample variance contribution')
#        plot(d1[:,0], d1[:,1]+d2[:,1], '-k', label='TOTAL')

#        legend(loc='best',frameon=False)

#        show()
